/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.analysys.presto.connector.hbase;

import com.analysys.presto.connector.hbase.utils.Utils;
import com.facebook.presto.common.type.Type;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.RecordCursor;
import com.facebook.presto.spi.RecordSet;
import io.airlift.log.Logger;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ResultScanner;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.filter.CompareFilter;
import org.apache.hadoop.hbase.filter.Filter;
import org.apache.hadoop.hbase.filter.FilterList;
import org.apache.hadoop.hbase.filter.FirstKeyOnlyFilter;
import org.apache.hadoop.hbase.filter.SingleColumnValueFilter;
import org.apache.hadoop.hbase.util.Bytes;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * HBase record set
 *
 * @author wupeng
 * @date 2019/01/29
 */
public class HBaseRecordSet implements RecordSet {

    private static final Logger log = Logger.get(HBaseRecordSet.class);

    private final List<HBaseColumnHandle> columnHandles;
    private final List<Type> columnTypes;
    private final HBaseSplit hBaseSplit;
    private final Connection connection;
    private Map<Integer, HBaseColumnHandle> fieldIndexMap = new HashMap<>();

    HBaseRecordSet(HBaseSplit split, List<ColumnHandle> columnHandles, HBaseClientManager clientManager) {
        this.hBaseSplit = Objects.requireNonNull(split, "split is null");
        Objects.requireNonNull(clientManager, "clientManager is null");
        Objects.requireNonNull(columnHandles, "column handles is null");
        this.columnHandles = columnHandles.stream().map(ch -> (HBaseColumnHandle) ch).collect(Collectors.toList());
        this.initFieldIndexMap(this.columnHandles);

        this.columnTypes = columnHandles.stream().map(ch -> ((HBaseColumnHandle) ch).getColumnType())
                .collect(Collectors.toList());

        this.connection = clientManager.createConnection();
    }

    @Override
    public List<Type> getColumnTypes() {
        return this.columnTypes;
    }

    @Override
    public RecordCursor cursor() {
        try {
            Table table = connection.getTable(TableName.valueOf(
                        hBaseSplit.getSchemaName() + ":" + hBaseSplit.getTableName()));
            // Check out if this is batch get
            if (Utils.isBatchGet(this.hBaseSplit.getConstraint(), hBaseSplit.getRowKeyName())) {
                return new HBaseGetRecordCursor(this.columnHandles,
                        this.hBaseSplit, this.fieldIndexMap, table, this.connection);
            } else {
                // client side region scanner
                long start = System.currentTimeMillis();
                Scan scan = getScanFromPrestoConstraint();
                ResultScanner resultScanner = table.getScanner(scan);
                log.info("scan hbase cost time: " + (System.currentTimeMillis() - start) + " ms");
                return new HBaseScanRecordCursor(this.columnHandles,
                        this.hBaseSplit,
                        resultScanner,
                        this.fieldIndexMap,
                        table,
                        this.connection);
            }
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
    }


    /**
     * Scan Filter
     *
     * @param condition
     * @return
     */
    private Filter getFilter(ConditionInfo condition) {
        CompareFilter.CompareOp operator;
        switch (condition.getOperator()) {
            case GT:
                operator = CompareFilter.CompareOp.GREATER;
                break;
            case GE:
                operator = CompareFilter.CompareOp.GREATER_OR_EQUAL;
                break;
            case LT:
                operator = CompareFilter.CompareOp.LESS;
                break;
            case LE:
                operator = CompareFilter.CompareOp.LESS_OR_EQUAL;
                break;
            default:
                operator = CompareFilter.CompareOp.EQUAL;
                break;
        }
        SingleColumnValueFilter f = new SingleColumnValueFilter(
                Bytes.toBytes(getFamilyByColumnName(condition.getColName(), columnHandles)),
                Bytes.toBytes(condition.getColName()), operator,
                condition.valueToBytes());
        f.setFilterIfMissing(true);
        return f;
    }

    private String getFamilyByColumnName(String columnName, List<HBaseColumnHandle> columns) {
        Objects.requireNonNull(columnName, "column name is null");
        HBaseColumnHandle column = columns.stream()
                .filter(col -> columnName.equals(col.getColumnName())).findAny().orElse(null);
        if (column != null) {
            return column.getFamily();
        } else {
            return "unknown_family";
        }
    }

    private Scan getScanFromPrestoConstraint() {
        Scan scan = new Scan();
        scan.setCaching(100);
        scan.setLoadColumnFamiliesOnDemand(true);
        scan.setCacheBlocks(true);
        scan.setLimit(hBaseSplit.getScanLimit());

        // Filter the exactly columns we want
        // for (HBaseColumnHandle hch : this.columnHandles) {
        this.columnHandles.forEach(hch -> {
            // 不是行转列的才加入，否则全部加入
            if(!this.hBaseSplit.isExplodeMode()) {
                if (this.hBaseSplit.getRowKeyName() == null) {
                    scan.addColumn(Bytes.toBytes(hch.getFamily()), Bytes.toBytes(hch.getColumnName()));
                } else {
                    if (!this.hBaseSplit.getRowKeyName().equals(hch.getColumnName())) {
                        scan.addColumn(
                                Bytes.toBytes(hch.getFamily()), Bytes.toBytes(hch.getColumnName()));
                    }
                }
            }
        });

        FilterList allFilters = new FilterList(FilterList.Operator.MUST_PASS_ALL);

        // ---------- Constraint push down ----------
        // This means user sql is like below:
        // select count(rowKey) / rowKey from table_xxx;
        // So we add FirstKeyOnlyFilter to return the first column to get the rowKey
        if (this.columnHandles.size() == 1
                && this.columnHandles.get(0).getColumnName().equals(this.hBaseSplit.getRowKeyName())) {
            allFilters.addFilter(new FirstKeyOnlyFilter());
            scan.setFilter(allFilters);
        } else {
            Map<String, List<ConditionInfo>> conditions = hBaseSplit.getConstraint().stream()
                    .collect(Collectors.groupingBy(ConditionInfo::getColName));
            // Here is what kind of condition presto can give to us:
            // 1.There can only be an 'and' relationship between different columns
            // 2.The same column can only be an 'or' relationship
            for (Map.Entry<String, List<ConditionInfo>> entry : conditions.entrySet()) {
                // Same column
                if (entry.getValue().size() > 1) {
                    List<Filter> columnFilterList = entry.getValue().stream().map(this::getFilter)
                            .collect(Collectors.toList());
                    FilterList columnFilter = new FilterList(FilterList.Operator.MUST_PASS_ONE, columnFilterList);
                    allFilters.addFilter(columnFilter);
                }
                // different columns
                else {
                    allFilters.addFilter(getFilter(entry.getValue().get(0)));
                }
            }
            if (hBaseSplit.getConstraint().size() >= 1) {
                scan.setFilter(allFilters);
            }
        }
        // ---------- Constraint push down finished ----------

        if (hBaseSplit.getStartRow() != null && hBaseSplit.getEndRow() != null) {
            scan.setStopRow(Bytes.toBytes(hBaseSplit.getEndRow()));
            scan.setStartRow(Bytes.toBytes(hBaseSplit.getStartRow()));
        }
        return scan;
    }

    private void initFieldIndexMap(List<HBaseColumnHandle> columnHandles) {
        columnHandles.forEach(hch -> fieldIndexMap.put(hch.getColumnName().hashCode(), hch));
    }

    HBaseSplit getHBaseSplit() {
        return hBaseSplit;
    }
}





















